package com.atguigu.guli.ai.service.impl;

import com.atguigu.common.core.domain.model.LoginUser;
import com.atguigu.guli.ai.constant.SystemConstant;
import com.atguigu.guli.ai.pojo.Chat;
import com.atguigu.guli.ai.pojo.Message;
import com.atguigu.guli.ai.service.AiService;
import com.atguigu.guli.ai.strategy.AiBean;
import com.atguigu.guli.ai.strategy.AiOperator;
import com.atguigu.guli.ai.util.FileUtil;
import com.atguigu.guli.ai.util.MongoUtil;
import com.atguigu.guli.ai.util.SnowflakeIdGenerator;
import com.atguigu.guli.ai.vo.ChatVo;
import com.atguigu.guli.ai.vo.MessageVo;
import com.atguigu.guli.ai.vo.QueryVo;
import com.atguigu.system.domain.ChatKnowledge;
import com.atguigu.system.domain.ChatProject;
import com.atguigu.system.mapper.ChatKnowledgeMapper;
import com.atguigu.system.mapper.ChatProjectMapper;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.data.domain.Sort;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.CollectionUtils;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

import static com.atguigu.common.utils.SecurityUtils.getLoginUser;

@Service
public class AiServiceImpl implements AiService, ApplicationContextAware {

    @Autowired
    private ChatKnowledgeMapper chatKnowledgeMapper;
    @Autowired
    private ChatProjectMapper chatProjectMapper;
    @Autowired
    private MongoTemplate mongoTemplate;

    private static final Map<String, AiOperator> MAP = new ConcurrentHashMap<>();

    public AiOperator getOperator(String type) {
        return MAP.get(type);
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        Map<String, Object> beans = applicationContext.getBeansWithAnnotation(AiBean.class);
        System.out.println("beans = " + beans);
        if (CollectionUtils.isEmpty(beans)) {
            return;
        }
        beans.values().forEach(bean->{
            System.out.println("bean = " + bean);
            AiBean aiBean = bean.getClass().getAnnotation(AiBean.class);
            System.out.println("aiBean = " + aiBean);
            MAP.put(aiBean.value(),(AiOperator) bean);
        });

    }

    @Transactional
    @Override
    public void upload(ChatKnowledge chatKnowledge, MultipartFile file) {
        String context = FileUtil.getContentFromText(file);
        chatKnowledge.setContent(context);
        chatKnowledge.setCreateTime(new Date());
        chatKnowledge.setFileName(file.getOriginalFilename());
        LoginUser loginUser = getLoginUser();
        String username = loginUser.getUsername();
        Long userId = loginUser.getUserId();
        chatKnowledge.setUserId(userId);
        chatKnowledge.setCreateBy(username);
        chatKnowledgeMapper.insertChatKnowledge(chatKnowledge);
        ChatProject chatProject = chatProjectMapper.selectChatProjectByProjectId(chatKnowledge.getProjectId());
        this.getOperator(chatProject.getType()).addDosc(chatKnowledge);
    }

    @Override
    public String creareChat(ChatVo chatVo) {
        Chat chat = new Chat();
        BeanUtils.copyProperties(chatVo, chat);
        SnowflakeIdGenerator snowflakeIdGenerator = new SnowflakeIdGenerator();
        long chatId = snowflakeIdGenerator.nextId();
        chat.setChatId(chatId);
        chat.setCreateTime(new Date());
        mongoTemplate.insert(chat, MongoUtil.getChatCollectionName(chatVo.getProjectId()));
        return chat.getChatId().toString();
    }

    @Override
    public List<Chat> getChat(ChatVo chatVo) {
        return mongoTemplate.find(Query.query(Criteria.where("projectId")
                        .is(chatVo.getProjectId())
                        .and("userId")
                        .is(chatVo.getUserId()))
                .with(Sort.by(Sort.Order.desc("createTime"))), Chat.class, MongoUtil.getChatCollectionName(chatVo.getProjectId()));
    }

    @Override
    public void updateChat(ChatVo chatVo) {
        mongoTemplate.updateFirst(Query.query(Criteria.where("_id")
                .is(chatVo.getChatId())),
                Update.update("title",chatVo.getTitle()),
                MongoUtil.getChatCollectionName(chatVo.getProjectId()));
    }

    @Override
    public Flux<String> chatStream(QueryVo queryVo) {
        // 记录用户的问题 将用户的问题保存到mongodb
        String msgCollectionName = MongoUtil.getMsgCollectionName(queryVo.getChatId());
        Long chatId = queryVo.getChatId();
        long msgId = new SnowflakeIdGenerator().nextId();
        Message message = new Message();
        message.setId(msgId);
        message.setType(0);
        message.setContent(queryVo.getMsg());
        message.setChatId(chatId);
        message.setCreateTime(new Date());
        mongoTemplate.insert(message, msgCollectionName);
        // 查询本地知识库
        ChatProject chatProject = chatProjectMapper.selectChatProjectByProjectId(queryVo.getProjectId());
        List<Document> docs = this.getOperator(chatProject.getType()).similaritySearch(queryVo);
        // 查询历史问答，作为联系上下文的提示
        List<Message> messages = mongoTemplate.find(Query.query(
                Criteria.where("chatId").is(chatId))
                .with(Sort.by(Sort.Order.asc("createTime")))
                , Message.class, msgCollectionName);

        // 组装上下文提示
        List<org.springframework.ai.chat.messages.Message> msgs = new ArrayList<>();
        if (!CollectionUtils.isEmpty(messages)){
            messages.forEach(messageAll -> {
                org.springframework.ai.chat.messages.Message msg = null;
                if (messageAll.getType().intValue() == 0){
                    msg = new UserMessage(messageAll.getContent());
                }else if (messageAll.getType().intValue() == 1){
                    msg = new AssistantMessage(messageAll.getContent());
                }
                msgs.add(msg);
            });
        }
        // 加上系统提示
        if (!CollectionUtils.isEmpty(docs)){
            docs.forEach(doc -> {
                msgs.add(new SystemMessage(doc.getContent()));
            });
        }
        Flux<String> result = this.getOperator(chatProject.getType()).chat_stream(msgs);

        return result;
    }

    @Override
    public void saveMsg(MessageVo messageVo) {
        String chatCollectionName = MongoUtil.getMsgCollectionName(messageVo.getChatId());
        long msgId = new SnowflakeIdGenerator().nextId();
        Message message = new Message();
        message.setId(msgId);
        message.setType(1);
        message.setChatId(messageVo.getChatId());
        message.setContent(messageVo.getContent());
        message.setCreateTime(new Date());
        mongoTemplate.insert(message, chatCollectionName);
    }

    @Override
    public List<Message> listMsg(Long chatId) {
        String msgCollectionName = MongoUtil.getMsgCollectionName(chatId);
        return mongoTemplate.find(Query.query(
                Criteria.where("chatId")
                        .is(chatId))
                .with(Sort.by(Sort.Order.asc("createTime"))),
                Message.class,
                msgCollectionName);
    }


}
